package com.coffee.mongodb;

import com.coffee.core.CoffeePage;
import org.springframework.data.domain.Sort;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.aggregation.*;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;


/**
 * @author: songkui
 * @since: 2022/6/23 11:29
 */
public class MochaMongoTemplate {

    private final MongoTemplate mongoTemplate;

    public MochaMongoTemplate(MongoTemplate mongoTemplate){
        this.mongoTemplate = mongoTemplate;
    }

    public MongoTemplate getMongoTemplate(){
        return mongoTemplate;
    }

    /**
     * 单表 分页 查询
     * @param mochaMongoQuery
     * @param entityClass
     * @param <T>
     * @return
     */
    public <T> CoffeePage<T> queryByPage(MochaMongoQuery<T> mochaMongoQuery, Class<T> entityClass){
        Query query = null == mochaMongoQuery.criteria ? new Query() : Query.query(mochaMongoQuery.criteria);

        if (null != mochaMongoQuery.fieldsExclude && mochaMongoQuery.fieldsExclude.length > 0){
            query.fields().exclude(mochaMongoQuery.fieldsExclude);
        }else if (null != mochaMongoQuery.fieldsInclude && mochaMongoQuery.fieldsInclude.length > 0){
            query.fields().include(mochaMongoQuery.fieldsInclude);
        }

        long total = mochaMongoQuery.searchCount ? mongoTemplate.count(query, entityClass) : mochaMongoQuery.total;
        if (StringUtils.hasLength(mochaMongoQuery.orderASC)){
            query.with(Sort.by(Sort.Direction.ASC, mochaMongoQuery.orderASC.split(",")));
        }
        if (StringUtils.hasLength(mochaMongoQuery.orderDESC)){
            query.with(Sort.by(Sort.Direction.DESC, mochaMongoQuery.orderDESC.split(",")));
        }
        query.skip(mochaMongoQuery.skip()).limit(mochaMongoQuery.limit());
        List<T> list = mongoTemplate.find(query, entityClass);
        return CoffeePage.of(list, total, mochaMongoQuery.size, mochaMongoQuery.current);
    }

    /**
     * 单表查询
     * @param mochaMongoQuery
     * @param entityClass
     * @param <T>
     * @return 返回符合条件的索有list;
     */
    public <T> List<T> queryByList(MochaMongoQuery<T> mochaMongoQuery, Class<T> entityClass){
        Query query = null == mochaMongoQuery.criteria ? new Query() : Query.query(mochaMongoQuery.criteria);

        if (null != mochaMongoQuery.fieldsExclude && mochaMongoQuery.fieldsExclude.length > 0){
            query.fields().exclude(mochaMongoQuery.fieldsExclude);
        }else if (null != mochaMongoQuery.fieldsInclude && mochaMongoQuery.fieldsInclude.length > 0){
            query.fields().include(mochaMongoQuery.fieldsInclude);
        }
        if (StringUtils.hasLength(mochaMongoQuery.orderASC)){
            query.with(Sort.by(Sort.Direction.ASC, mochaMongoQuery.orderASC.split(",")));
        }
        if (StringUtils.hasLength(mochaMongoQuery.orderDESC)){
            query.with(Sort.by(Sort.Direction.DESC, mochaMongoQuery.orderDESC.split(",")));
        }
        return mongoTemplate.find(query, entityClass);
    }

    /**
     * 联合查询
     * @param mochaMongoQuery
     * @param entityClass
     * @param <T>
     * @return
     */
    public <T> CoffeePage<T> joinQueryByPage(MochaMongoQuery<T> mochaMongoQuery, Class<T> entityClass)throws MochaMongoException {
        if (null == mochaMongoQuery.joins || mochaMongoQuery.joins.length == 0){
            throw new MochaMongoException("Mongo join query joins is null!");
        }
        if (!StringUtils.hasLength(mochaMongoQuery.master)){
            throw new MochaMongoException("Mongo join query master collectionName  is null!");
        }

        List<AggregationOperation> operations = new ArrayList<>();

        addLookupOperation(mochaMongoQuery, operations);

        if (null != mochaMongoQuery.criteria){
            operations.add(Aggregation.match(mochaMongoQuery.criteria));
        }

        long total = totalQuery(mochaMongoQuery, operations);
        if (total == 0L){
            return CoffeePage.of(0L, mochaMongoQuery.size, mochaMongoQuery.current);
        }

        addProjectionOperation(mochaMongoQuery, operations);

        addSortOperation(mochaMongoQuery, operations);
        //分页
        operations.add(Aggregation.skip(mochaMongoQuery.skip()));
        operations.add(Aggregation.limit(mochaMongoQuery.limit()));

        Aggregation aggregation = Aggregation.newAggregation(operations);
        AggregationResults<T> results =  mongoTemplate.aggregate(aggregation, mochaMongoQuery.master, entityClass);
        return CoffeePage.of(results.getMappedResults(), total, mochaMongoQuery.size, mochaMongoQuery.current);
    }

    /**
     * 联合查询 返回 list;
     * @param mochaMongoQuery
     * @param entityClass
     * @param <T>
     * @return
     * @throws MochaMongoException
     */
    public <T> List<T> joinQueryByList(MochaMongoQuery<T> mochaMongoQuery, Class<T> entityClass)throws MochaMongoException {
        if (null == mochaMongoQuery.joins || mochaMongoQuery.joins.length == 0){
            throw new MochaMongoException("Mongo join query joins is null!");
        }
        if (!StringUtils.hasLength(mochaMongoQuery.master)){
            throw new MochaMongoException("Mongo join query master collectionName  is null!");
        }

        List<AggregationOperation> operations = new ArrayList<>();

        addLookupOperation(mochaMongoQuery, operations);

        if (null != mochaMongoQuery.criteria){
            operations.add(Aggregation.match(mochaMongoQuery.criteria));
        }

        addProjectionOperation(mochaMongoQuery, operations);

        addSortOperation(mochaMongoQuery, operations);

        Aggregation aggregation = Aggregation.newAggregation(operations);
        AggregationResults<T> results =  mongoTemplate.aggregate(aggregation, mochaMongoQuery.master, entityClass);
        return results.getMappedResults();
    }

    private <T> void  addLookupOperation(MochaMongoQuery<T> mochaMongoQuery, List<AggregationOperation> operations){
        for (MochaMongoQuery.Join join : mochaMongoQuery.joins){
            LookupOperation lookupOne = LookupOperation.newLookup()
                    .from(join.from)
                    .localField(join.localField)
                    .foreignField(join.foreignField)
                    .as(join.as);
            operations.add(lookupOne);
            if (join.unwind){
                operations.add(Aggregation.unwind(join.as)); // 返回对象；  否则 返回 list
            }
        }
    }

    private <T>long totalQuery(MochaMongoQuery<T> mochaMongoQuery, List<AggregationOperation> operations){
        if (mochaMongoQuery.searchCount || mochaMongoQuery.total == 0){
            List<AggregationOperation> countOperations = new ArrayList<>();
            countOperations.addAll(operations);
            countOperations.add(Aggregation.count().as("count"));
            AggregationResults<Map> countMap = mongoTemplate.aggregate(Aggregation.newAggregation(countOperations), mochaMongoQuery.master, Map.class);
            List<Map> countList =  countMap.getMappedResults();
            if (null == countList || countList.isEmpty()){
                return 0L;
            }
            return Long.parseLong(String.valueOf(countList.get(0).get("count")));
        }

        return mochaMongoQuery.total;
    }

    /**
     * 指定了字段那么 as 对象里面就不会赋值； 会封装到对应字段
     * @param mochaMongoQuery
     * @param operations
     * @param <T>
     */
    private <T> void addProjectionOperation(MochaMongoQuery<T> mochaMongoQuery, List<AggregationOperation> operations){
        if (null != mochaMongoQuery.fieldsInclude && mochaMongoQuery.fieldsInclude.length > 0){
            ProjectionOperation projectionOperation;
            List<String> notAsStrs = Arrays.stream(mochaMongoQuery.fieldsInclude).filter(s -> !s.contains(" as ")).collect(Collectors.toList());
            if (null != notAsStrs && !notAsStrs.isEmpty()){
                String[] in =  notAsStrs.toArray(new String[notAsStrs.size()]);
                projectionOperation = Aggregation.project(in);
            }else {
                projectionOperation = new ProjectionOperation();
            }

            for (String include : mochaMongoQuery.fieldsInclude){
                if (include.contains(" as ")){
                    String[] as = include.split(" as ");
                    projectionOperation = projectionOperation.and(as[0].trim()).as(as[1]);
                }
            }
            operations.add(projectionOperation);
        }else if (null != mochaMongoQuery.fieldsExclude && mochaMongoQuery.fieldsExclude.length > 0){
            ProjectionOperation projectionOperation = new ProjectionOperation();
            operations.add(projectionOperation.andExclude(mochaMongoQuery.fieldsExclude));
        }

    }

    private <T> void addSortOperation(MochaMongoQuery<T> mochaMongoQuery, List<AggregationOperation> operations){
        if (StringUtils.hasLength(mochaMongoQuery.orderASC)){
            operations.add(Aggregation.sort(Sort.Direction.ASC, mochaMongoQuery.orderASC.split(",")));
        }
        if (StringUtils.hasLength(mochaMongoQuery.orderDESC)){
            operations.add(Aggregation.sort(Sort.Direction.DESC, mochaMongoQuery.orderDESC.split(",")));
        }
    }


}
